package com.mixed.mixedserver.controller;

import com.mixed.mixedserver.ai.prompt.LawPrompt;
import com.mixed.mixedserver.service.ai.ChatAssistant;
import com.mixed.mixedserver.service.ai.ChatMemoryAssistant;
import com.mixed.mixedserver.service.ai.LawChatAssistant;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ImageContent;
import dev.langchain4j.data.message.TextContent;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.model.input.Prompt;
import dev.langchain4j.model.input.PromptTemplate;
import jakarta.annotation.Resource;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.Base64;
import java.util.Map;

@RestController
@RequestMapping("/ai")
public class AiChatController {

    @Resource(name = "QWen")
    private ChatModel chatModel;

    @Resource(name = "DeepSeek")
    private ChatModel deepSeek;

    @Resource
    private ChatAssistant chatAssistant;

    @Resource(name = "QWenImageModel")
    private ChatModel qWenImageModel;

    @Value("classpath:static/img.png")
    private org.springframework.core.io.Resource img;

    @Resource(name = "chatAssistantFlux")
    private ChatAssistant chatAssistantFlux;

    @Resource
    private StreamingChatModel streamingChatModel;

    @Resource(name = "chatTokenWindowChatMemory")
    private ChatMemoryAssistant chatMemoryAssistant;

    @GetMapping("/chatWithQWen")
    public String chat(@RequestParam(defaultValue = "你是谁") String question) {
        String chat = chatModel.chat(question);
        return chat;
    }

    @GetMapping("/chatWithDeepSeek")
    public String chatWithDeepSeek(@RequestParam(defaultValue = "你是谁") String question) {
        String chat = deepSeek.chat(question);
        return chat;
    }

    @GetMapping("/chatWithAssistant")
    public String chatWithAssistant(@RequestParam(value = "prompt", defaultValue = "你是谁") String prompt) {
        String chat = chatAssistant.chat(prompt);
        return chat;
    }

    @GetMapping("/chatWithImageQwen")
    public String chatWithImageQWen(@RequestParam(value = "prompt", defaultValue = "理解图片内容") String prompt) throws IOException {
        String result = "";
        // 1.图片转码
        byte[] contentAsByteArray = img.getContentAsByteArray();
        String encode = Base64.getEncoder().encodeToString(contentAsByteArray);
        // 2.提示词指定
        UserMessage userMessage = UserMessage.from(TextContent.from("请用中文解释图片内容，并给出图片的描述"), ImageContent.from(encode, "image/png"));
        // 3.调用模型
        ChatResponse chat = qWenImageModel.chat(userMessage);
        // 4.返回结果
        AiMessage aiMessage = chat.aiMessage();
        result = aiMessage.text();
        return result;
    }

    @GetMapping("/chatStream/chat")
    public Flux<String> chatStream(@RequestParam("prompt") String prompt) {
       return Flux.create(sink -> {
            streamingChatModel.chat(prompt, new StreamingChatResponseHandler() {
                @Override
                public void onPartialResponse(String s) {
                    sink.next(s);
                }

                @Override
                public void onCompleteResponse(ChatResponse chatResponse) {
                    sink.complete();
                }

                @Override
                public void onError(Throwable throwable) {
                    sink.error(throwable);
                }
            });
        });
    }

    @GetMapping("/chatStream/chatAssistantFlux")
    public Flux<String> chatStreamAssistantFlux(@RequestParam("prompt") String prompt) {
        return chatAssistantFlux.chatStream(prompt);
    }

    @GetMapping("/chatTokenWindowChatMemory")
    public String chatTokenWindowChatMemory(@RequestParam("prompt") String prompt,@RequestParam("memoryId") String memoryId) {
        return chatMemoryAssistant.chat(memoryId, prompt);
    }

    @Resource
    private LawChatAssistant lawChatAssistant;

    @GetMapping("/lawChatAssistant")
    public String lawChatAssistant(@RequestParam("prompt") String prompt) {
        return lawChatAssistant.chat( prompt,200,"markdown");
    }

    @GetMapping("/lawChatAssistant2")
    public String lawChatAssistant2(@RequestParam("question") String question) {
        return lawChatAssistant.chat(new LawPrompt("markdown",question));
    }

    @GetMapping("/chat/promptTemplate")
    public String chatWithPromptTemplate(@RequestParam("question") String question) {
        String result = "";
        // 1.构造模板
        PromptTemplate promptTemplate = PromptTemplate.from("请用中文回答以下问题:{{question}}");
        // 2.生成prompt
        Prompt apply = promptTemplate.apply(Map.of("it", "编程学习助手", "question", question));
        // 3.提示词变为UserMessage
        UserMessage userMessage = apply.toUserMessage();
        // 4.调用大模型
        ChatResponse chat = chatModel.chat(userMessage);
        result = chat.aiMessage().text();
        return result;
    }
}
